# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: daniel

import collections
import os
import matplotlib.pyplot as plt
import json
from absl import app, flags


FLAGS = flags.FLAGS


class DebugInfoAnalyzer(object):
  def __init__(self):
    error_stats = collections.defaultdict(lambda: 0)
    error_stats["ORDER_CONTEXT_DEAD_ORDER_RESURRECTED"] = 0
    error_stats['SUBMIT_PENDING_TIMEOUT'] = 0
    error_stats['CANCEL_PENDING_TIMEOUT'] = 0
    error_stats['REMVOE_ZOMBIE_ORDER'] = 0
    error_stats["ORDER_CONTEXT_DEAD_ORDER_APPEARS_ALIVE"] = 0
    # error_stats["CANCELED_ORDER_HAS_NONZERO_WORKING_ORDER"] = 0
    self._error_stats = error_stats

  def _reset_error_stats(self):
    error_stats = self._error_stats
    error_stats["ORDER_CONTEXT_DEAD_ORDER_RESURRECTED"] = 0
    error_stats['SUBMIT_PENDING_TIMEOUT'] = 0
    error_stats['CANCEL_PENDING_TIMEOUT'] = 0
    error_stats['REMVOE_ZOMBIE_ORDER'] = 0
    error_stats["ORDER_CONTEXT_DEAD_ORDER_APPEARS_ALIVE"] = 0
    # error_stats["CANCELED_ORDER_HAS_NONZERO_WORKING_ORDER"] = 0

  def filter_debug_info(self, records):
    ret = []
    for record in records:
      if record['type'] != 'OG_LOG':
        continue
      if 'event' not in record['og_log']:
        continue
      event = record['og_log']['event']
      if event['type'] != 'DEBUG_INFO':
        continue
      ret.append(record)
    return ret

  def get_error_stats(self, info):
    error_stats = self._error_stats
    error_code = info.get('error_code')
    if error_code is None:
      return False

    if "ORDER_CONTEXT_DEAD_ORDER_RESURRECTED" in error_code:
      error_stats["ORDER_CONTEXT_DEAD_ORDER_RESURRECTED"] += 1

    if 'submit_pending_timeout' in info:
      error_stats['SUBMIT_PENDING_TIMEOUT'] += 1

    if 'SetTimeoutCancelRejected' in error_code:
      error_stats['CANCEL_PENDING_TIMEOUT'] += 1

    if 'max_query_alive_time' in info:
      error_stats['REMVOE_ZOMBIE_ORDER'] += 1

    if "ORDER_CONTEXT_DEAD_ORDER_APPEARS_ALIVE" in error_code:
      error_stats["ORDER_CONTEXT_DEAD_ORDER_APPEARS_ALIVE"] += 1

    # if "CANCELED_ORDER_HAS_NONZERO_WORKING_ORDER" in error_code:
    #  error_stats["CANCELED_ORDER_HAS_NONZERO_WORKING_ORDER"] += 1

  def retrieve_debug_info_record(self, records):
    for record in records:
      # account_request = record['og_log']['account_request']
      debug_info = json.loads(record['og_log']['event']['tag'])
      if 'error_code' in debug_info and isinstance(debug_info['error_code'], dict):
        continue

      for key, info_list in debug_info.items():
        if key == "og_version":
          continue

        if isinstance(info_list, list):
          for info in info_list:
            res = self.get_error_stats(info)
            if res is False:
              print('format error: ', info_list)
        else:
          print('should be a list: ', record)

  def generate_debug_info_report(self,
                                 records,
                                 strategy_name,
                                 machine,
                                 start_dt,
                                 end_dt,
                                 output_dir,
                                 local_run):
    self._reset_error_stats()
    debug_info_records = self.filter_debug_info(records)
    # ignore strategy with zero http conn timeout
    if len(debug_info_records) == 0:
      return
    trade_date = start_dt.strftime('%Y%m%d')
    start_dt_str = start_dt.strftime('%Y%m%dT%H%M%S')
    end_dt_str = end_dt.strftime('%Y%m%dT%H%M%S')

    self.retrieve_debug_info_record(debug_info_records)
    name = strategy_name

    x = sorted(self._error_stats.keys())
    y = []
    for key in x:
      y.append(self._error_stats[key])
    x_pos = [i for i, _ in enumerate(x)]

    if sum(y) == 0:
      return

    plt.rcParams.update({'font.size': 12})
    fig = plt.figure(figsize=(10, 6), dpi=200)
    plt.bar(x_pos, y, color='green')
    plt.xlabel("error types")
    plt.ylabel("# of errors")
    plt.title("order state abnormality (%s)" % name)
    plt.xticks(range(len(x)), x, rotation=270)
    plt.grid(True)
    plt.legend()
    fig.tight_layout()

    output_path = f"{output_dir}/{trade_date}/order_state_abnormality/{strategy_name}/{machine}"
    if not os.path.exists(output_path):
      os.makedirs(output_path)

    output_file = f"{start_dt_str}-{end_dt_str}.svg"
    print(f"write to {output_path}/{output_file}")
    fig.savefig(f"{output_path}/{output_file}", dpi=fig.dpi)
    plt.close(fig)


def main():
  # TODO(daniel):
  print("@daniel: to impliment")


if __name__ == '__main__':
  flags.DEFINE_string('root_dir',
                      '/remote/iosg/strat-1/buckets/log.raw.coin/live/strat_proto_log',
                      'root dir specify')
  flags.DEFINE_string('output_dir',
                      '/remote/iosg/strat-1/buckets/analysis.derived.coin/live/plot/og_log_plot',
                      'output_dir.')
  flags.DEFINE_string('tmp_dir',
                      '/tmp',
                      'tmp dir for slurm run')
  flags.DEFINE_string('kafka_config_filename', None, 'kafka config')
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')
  flags.DEFINE_string('machine', None, 'Instance machine name.')
  flags.DEFINE_string('strategy_name', None, 'Strategy name.')
  flags.DEFINE_string('start_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('end_time', None, '%Y%m%dT%H%M%S')
  flags.DEFINE_string('symbol', None, '')
  flags.DEFINE_string('log_type', None, '')
  flags.DEFINE_bool('local_run', False, "")
  flags.DEFINE_bool('slurm_run', False, "")
  flags.DEFINE_bool('generate_summary_report', False, "generate summary report only")
  app.run(main)
